# 数据管理器模块 - 用于协调各数据存储系统之间的数据流转

import logging
import pandas as pd
from typing import Dict, List, Any, Optional, Union
from datetime import datetime, timedelta

# 导入各存储系统客户端
from .redis_client import RedisClient
from .postgres_client import PostgresClient
from .parquet_storage import ParquetStorage
from src.utils.database import REDIS_CONFIG, POSTGRES_CONFIG, ARCHIVE_DIR
from src.utils.schemas import STOCK_PRICES_SCHEMA, TECHNICAL_INDICATORS_SCHEMA, STOCK_LIST_SCHEMA
from src.utils.logger_manager import log, err

logger = logging.getLogger(__name__)

class DataManager:
    _instance = None
    
    def __new__(cls):
        if cls._instance is None:
            cls._instance = super().__new__(cls)
            cls._instance._initialized = False
        return cls._instance
    
    def __init__(self):
        if not self._initialized:
            self.redis = RedisClient(**REDIS_CONFIG)
            self.postgres = PostgresClient(**POSTGRES_CONFIG)
            self.archive_dir = ARCHIVE_DIR
            self._initialized = True
    
    @classmethod
    def get_instance(cls):
        if cls._instance is None:
            cls._instance = cls()
        return cls._instance

    """
    数据管理器类，用于协调各数据存储系统之间的数据流转
    实现完整的数据生命周期管理：
    Redis(短期实时数据) -> PostgreSQL(历史数据) -> Parquet(冷数据归档)
    """
    def update_stock_list(self, new_stocks: Union[List[Dict], pd.DataFrame]):
        """
        更新股票列表，将不存在的股票新增到 stock_list 表。

        Args:
            new_stocks (Union[List[Dict], pd.DataFrame]): 包含股票代码和名称的数组或 DataFrame。
        """
        try:
            # 如果传入的是数组，转换为 DataFrame
            if isinstance(new_stocks, list):
                new_stocks = pd.DataFrame(new_stocks)

            if new_stocks.empty:
                logger.warning("新股票列表为空，跳过更新操作")
                return

            # 查询现有的股票列表
            existing_stocks = self.postgres.query_table(
                table_name='stock_list',
                columns=['code']
            )
            existing_codes = set(existing_stocks['code']) if not existing_stocks.empty else set()

            # 筛选出新增的股票
            new_stocks_to_add = new_stocks[~new_stocks['code'].isin(existing_codes)]

            if new_stocks_to_add.empty:
                logger.info("没有新增的股票需要更新")
                return

            # 插入新增的股票到 stock_list 表
            self.postgres.store_data('stock_list', new_stocks_to_add)
            logger.info(f"成功新增 {len(new_stocks_to_add)} 条股票到 stock_list 表")
        except Exception as e:
            logger.error(f"更新股票列表失败: {str(e)}")
    
    def query_stock_prices(self, symbol: str, start_date: Optional[str] = None, 
                           end_date: Optional[str] = None, con: Optional[Dict[str, Any]] = None, limit: Optional[int] = None, 
                           order_by: Optional[str] = 'start_date ASC') -> pd.DataFrame:
        """
        查找股票行情数据。

        Args:
            symbol (str): 股票代码。
            start_date (str, optional): 开始日期，格式'YYYY-MM-DD'。
            end_date (str, optional): 结束日期，格式'YYYY-MM-DD'。
            limit (int, optional): 限制返回的行数。
            order_by (str, optional): 排序条件，默认按日期升序。

        Returns:
            pd.DataFrame: 查询结果。
        """
        try:
            # 构建查询条件
            conditions = {'symbol = ': symbol}
            if start_date:
                conditions['start_date >='] = start_date
            if end_date:
                conditions['end_date <='] = end_date
            if con:
                for key, value in con.items():
                    conditions[key] = value

            # 调用PostgresClient的query_table方法
            return self.postgres.query_table(
                table_name='stock_prices',
                conditions=conditions,
                order_by=order_by,
                limit=limit
            )
        except Exception as e:
            logger.error(f"查询股票行情数据失败: {str(e)}")
            return pd.DataFrame()
    
    def store_stock_prices(self, data: Union[Dict, pd.DataFrame]):
        """存储股票价格数据"""
        try:
            table_name = 'stock_prices'

            # 检查数据类型并转换为 DataFrame
            if isinstance(data, dict):
                # 将字典转换为 DataFrame
                data = pd.DataFrame.from_records([data])

            if not isinstance(data, pd.DataFrame):
                raise ValueError("传入的数据必须是字典或DataFrame格式")

            if data.empty:
                logger.warning("数据为空，无法存储股票价格数据")
                return
            
            # 确保 start_date 和 end_date 的类型一致
            if 'start_date' in data.columns and not pd.api.types.is_datetime64_any_dtype(data['start_date']):
                data['start_date'] = pd.to_datetime(data['start_date'], errors='coerce').dt.date
            if 'end_date' in data.columns and not pd.api.types.is_datetime64_any_dtype(data['end_date']):
                data['end_date'] = pd.to_datetime(data['end_date'], errors='coerce').dt.date

            # 删除 'date' 列
            if 'date' in data.columns:
                data = data.drop(columns=['date'])
                
            if (not 'id' in data.columns) or len(data) > 10:
                self.postgres.store_data(table_name, data)
            else:
                for _, row in data.iterrows():
                    row_dict = row.to_dict()  # 确保每一行是字典类型
                    existing_data = self.postgres.query_table(
                        table_name=table_name,
                        conditions={'id =': row_dict['id'], 'start_date =': row_dict['start_date']},
                    )
                    if existing_data.empty:
                        # 如果不存在，则新增数据
                        self.postgres.store_data(table_name, pd.DataFrame([row_dict]))
                    else:
                        # 删除 id 和 start_date 字段，避免更新时修改这些字段
                        row_dict.pop('id', None)
                        row_dict.pop('start_date', None)
                        self.postgres.update_data(
                            table_name=table_name,
                            data=row_dict,
                            conditions={
                                'id =': int(existing_data['id'].iloc[-1]),  # 转换为 Python 原生 int 类型
                                'start_date =': str(existing_data['start_date'].iloc[-1])  # 转换为字符串
                            }
                        )
            log(f'股票 {data['symbol'].iloc[-1]}-{data['period'].iloc[-1]}历史行情数据存储成功，共{len(data)}条数据，最新日期为 {data['end_date'].iloc[-1]}')
        except Exception as e:
            err(f'股票 {data['symbol'].iloc[-1]}-{data['period'].iloc[-1]}历史行情数据存储失败')

    
    def store_technical_indicators(self, data: pd.DataFrame):
        """
        存储技术指标数据到技术指标表。

        Args:
            data (pd.DataFrame): 包含技术指标数据的 DataFrame。
        """
        try:
            table_name = 'technical_indicators'
            if data.empty:
                logger.warning("技术指标数据为空，无法存储")
                return

            # 筛选出技术指标表的字段，去除 created_at 字段
            valid_columns = set(TECHNICAL_INDICATORS_SCHEMA['columns'].keys()) - {'created_at'}
            data = data[[col for col in data.columns if col in valid_columns]]

            if len(data) > 10:
                self.postgres.store_data(table_name, data)
            else:
                for _, row in data.iterrows():
                    row_dict = row.to_dict()  # 确保每一行是字典类型
                    existing_data = self.postgres.query_table(
                        table_name=table_name,
                        conditions={'price_id =': row_dict['price_id'], 'start_date =': row_dict['start_date']},
                    )
                    if existing_data.empty:
                        # 如果不存在，则新增数据
                        self.postgres.store_data(table_name, pd.DataFrame([row_dict]))
                    else:
                        row_dict.pop('price_id', None)
                        row_dict.pop('id', None)
                        row_dict.pop('start_date', None)
                        self.postgres.update_data(
                            table_name=table_name,
                            data=row_dict,
                            conditions={'price_id =':int(existing_data['id'].iloc[-1]), 
                                        'start_date =': existing_data['start_date'].iloc[-1]}
                        )            
            log(f'技术指标存储成功，共{len(data)}条数据')
        except Exception as e:
            err(f'技术指标存储失败， {data['price_id'].iloc[-1]}-{data['start_date'].iloc[-1]}')
    
    def run_daily_tasks(self):
        """运行每日数据管理任务"""
        try:
            now = datetime.now()
            # 只在每天收盘后运行(16:00后)
            if now.hour < 16:
                log("未到收盘时间，跳过数据同步任务")
                return False

            # 检查是否为交易日
            trade_date = now.strftime('%Y%m%d')
            trading_days = self.redis.get('market:trading_days')
            if not trading_days or trade_date not in trading_days.split(','):
                log("非交易日，跳过数据同步任务")
                return False
            log("每日数据管理任务执行完成")

            log("开始执行每日数据管理任务")
            
            return True
        except Exception as e:
            logger.error(f"执行每日数据管理任务失败: {str(e)}")
            return False
        
    def archive_postgres_to_parquet(self, table_name: str, start_date: str = None, end_date: str = None):
        """
        将PostgreSQL中的数据归档到Parquet文件
        
        Args:
            table_name: 表名
            start_date: 开始日期，格式'YYYYMMDD'
            end_date: 结束日期，格式'YYYYMMDD'
        """
        try:
            # 如果未指定日期范围，默认归档最近一个月的数据
            if not start_date:
                start_date = (datetime.now() - timedelta(days=30)).strftime('%Y%m%d')
            if not end_date:
                end_date = datetime.now().strftime('%Y%m%d')
            
            # 归档数据
            # self.parquet.archive_from_postgres(self.postgres, table_name, start_date, end_date)
            
            logger.info(f"成功将PostgreSQL的{table_name}表数据归档到Parquet文件，日期范围: {start_date} 至 {end_date}")
            return True
        except Exception as e:
            logger.error(f"将PostgreSQL数据归档到Parquet文件失败: {str(e)}")
            return False
    
    def close(self):
        """
        关闭所有数据库连接
        """
        try:
            self.redis.close()
            self.postgres.close()
            logger.info("所有数据库连接已关闭")
        except Exception as e:
            logger.error(f"关闭数据库连接失败: {str(e)}")
    
    def __del__(self):
        """
        析构函数，确保所有连接被关闭
        """
        self.close()